import cv2
import numpy as np
import torch
import torch.utils.data as data

from .data_augment import mosaic


class WiderFaceDetection(data.Dataset):

    def __init__(self, txt_path, preproc=None, img_dim=640, cache=False):
        self.preproc = preproc
        self.imgs_path = []
        self.words = []
        self.img_dim = img_dim
        self.cache = cache
        self.cache_imgs = []
        f = open(txt_path, 'r')
        lines = f.readlines()
        isFirst = True
        labels = []
        for line in lines:
            line = line.rstrip()
            if line.startswith('#'):
                if isFirst is True:
                    isFirst = False
                else:
                    labels_copy = labels.copy()
                    self.words.append(labels_copy)
                    labels.clear()
                path = line[2:]
                path = txt_path.replace('label.txt', 'images/') + path
                self.imgs_path.append(path)
            else:
                line = line.split(' ')
                label = [float(x) for x in line]
                labels.append(label)

        self.words.append(labels)
        if self.cache:
            self.cache_imgs = [cv2.imread(self.imgs_path[index]) for index in range(len(self.imgs_path))]

    def __len__(self):
        return len(self.imgs_path)

    def __getitem__(self, index):
        if 0 == np.random.randint(0, 5):
            if self.cache:
                img, labels = mosaic(self.cache_imgs, self.words, index, self.img_dim)
            else:
                img, labels = mosaic(self.imgs_path, self.words, index, self.img_dim)
            # 将score < 0.2的标志位置为-1
            pos_index = np.where(labels[:, -1] >= 0.5)
            neg_index = np.where(labels[:, -1] < 0.5)
            labels[pos_index, -1] = 1
            labels[neg_index, -1] = -1
            # 将关键点标志位==-1的行扔掉
            keypoints_visible = labels[:, [6, 9, 12, 15, 18]]
            labels = labels[np.amin(keypoints_visible, axis=1) != -1]
            labels = labels[:, [0, 1, 2, 3, 4, 5, 7, 8, 10, 11, 13, 14, 16, 17, 19]]
            # 计算w h
            labels_wh = np.c_[labels[:, 2] - labels[:, 0], labels[:, 3] - labels[:, 1]]
            # 去除掉w, h小于15的框
            labels = labels[np.all(labels_wh[:, :] >= 15.0, axis=1)]
            labels[:, 0:-1:2] /= img.shape[2]
            labels[:, 1:-1:2] /= img.shape[1]

            return torch.from_numpy(img), labels
        else:
            if self.cache:
                img = self.cache_imgs[index]
            else:
                img = cv2.imread(self.imgs_path[index])
            height, width, _ = img.shape

            labels = self.words[index]
            annotations = np.zeros((0, 15))
            if len(labels) == 0:
                return annotations
            for idx, label in enumerate(labels):
                annotation = np.zeros((1, 15))
                # bbox
                annotation[0, 0] = label[0]  # x1
                annotation[0, 1] = label[1]  # y1
                annotation[0, 2] = label[0] + label[2]  # x2
                annotation[0, 3] = label[1] + label[3]  # y2

                # landmarks
                # xy-> 5 points
                annotation[:, 4:14] = np.array(label[4:-1]).reshape(-1, 3)[:, :2].reshape(-1)
                scores = np.array(label[4:-1]).reshape(-1, 3)[:, -1:].reshape(-1)
                clarity = label[-1]
                annotation[0, 14] = 1
                if clarity < 0.4 or len(np.where(scores < 0)[0]):
                    annotation[0, 14] = -1

                annotations = np.append(annotations, annotation, axis=0)
            target = np.array(annotations)

            if self.preproc is not None:
                img, target = self.preproc(img, target)

            target = target[np.all(abs(target[:, :]) <= 1.0, axis=1)]

            return torch.from_numpy(img), target


def detection_collate(batch):
    """Custom collate fn for dealing with batches of images that have a different
    number of associated object annotations (bounding boxes).

    Arguments:
        batch: (tuple) A tuple of tensor images and lists of annotations

    Return:
        A tuple containing:
            1) (tensor) batch of images stacked on their 0 dim
            2) (list of tensors) annotations for a given image are stacked on 0 dim
    """
    targets = []
    imgs = []
    for _, sample in enumerate(batch):
        for _, tup in enumerate(sample):
            if torch.is_tensor(tup):
                imgs.append(tup)
            elif isinstance(tup, type(np.empty(0))):
                annos = torch.from_numpy(tup).float()
                targets.append(annos)

    return torch.stack(imgs, 0), targets
